<?php
/**
 * Project: AIGC
 * Author: moufer <moufer@163.com>
 * Date: 2024/7/6 23:53
 */

namespace module\ai\library\channel\llm;

use Generator;
use module\ai\library\channel\abstract\Completion;
use module\ai\library\channel\abstract\Llm;
use module\ai\library\channel\completion\TencentCompletion;
use module\ai\library\channel\exception\CompletionException;
use module\ai\library\channel\interface\CompletionInterface;
use mof\exception\LogicException;
use TencentCloud\Common\Credential;
use TencentCloud\Common\Profile\ClientProfile;
use TencentCloud\Common\Profile\HttpProfile;
use TencentCloud\Hunyuan\V20230901\HunyuanClient;
use TencentCloud\Hunyuan\V20230901\Models\ChatCompletionsRequest;
use TencentCloud\Hunyuan\V20230901\Models\ChatCompletionsResponse;
use think\helper\Str;

class TencentHunyuan extends Llm
{
    protected string $host    = 'hunyuan.tencentcloudapi.com';
    protected string $baseUrl = 'https://hunyuan.tencentcloudapi.com';

    /**
     * @inheritdoc
     */
    public static function versions(): array
    {
        return [
            'hunyuan-lite'          => 'hunyuan-lite',
            'hunyuan-standard'      => 'hunyuan-standard-32k',
            'hunyuan-standard-256K' => 'hunyuan-standard-256k',
            'hunyuan-pro'           => 'hunyuan-pro',
            'hunyuan-role'          => 'hunyuan-role',
            'hunyuan-code'          => 'hunyuan-code',
            'hunyuan-functioncall'  => 'hunyuan-functioncall',
            'hunyuan-vision'        => 'hunyuan-vision',
        ];
    }

    /**
     * @inheritdoc
     */
    protected function setPayload(array $config, bool $override = false): void
    {
        parent::setPayload($config, $override);

        $keys = [
            'top_p', 'temperature', 'enable_enhancement',
        ];
        foreach ($keys as $key) {
            if (isset($config[$key])) {
                $this->payload[Str::studly($key)] = $config[$key];
            }
        }

        if ($this->payload['Model'] === 'hunyuan-functioncall') {
            if (empty($config['tools']) || empty($config['custom_tool'])) {
                throw new LogicException('混元大模型工具未配置完整');
            }
            $this->payload['Tools'] = $config['tools'];
            $this->payload['ToolChoice'] = $config['tool_choice'] ?? 'auto';
            $this->payload['CustomTool'] = $config['custom_tool'];
        }

        $defaultValues = [
            'action'  => 'ChatCompletion',
            'version' => '2023-09-01',
        ];
        foreach ($defaultValues as $key => $value) {
            $_key = Str::snake($key);
            if (empty($this->payload[$_key])) {
                $this->payload[$_key] = $value;
            }
        }
    }

    public function withModel(string $model): static
    {
        $this->payload['Model'] = $model;
        return $this;
    }

    /**
     * @inheritdoc
     */
    public function withMessage(string $message, string $role = 'user', ?array $extra = null, bool $overload = false): static
    {
        $message = [
            'Role'    => $role,
            'Content' => $message,
        ];
        if ($extra) {
            $message = array_merge($message, $extra);
        }
        if ($overload) {
            $this->payload['Messages'] = [];
        }
        $this->payload['Messages'][] = $message;
        return $this;
    }

    public function completion(): Completion
    {
        $response = $this->getCompletionsResponse();
        return $this->getCompletionStruct(json_decode($response->toJsonString(), true));
    }

    public function completionStream(): Generator
    {
        yield $this->getCompletionsResponse(function ($chunk) {
            $data = json_decode($chunk);
            yield $this->getCompletionStruct($data, true);
        });
    }

    /**
     * @inheritdoc
     */
    protected function getCompletionStruct(array $data, bool $isStream = false): ?CompletionInterface
    {
        return new TencentCompletion($data, $isStream);
    }

    protected function getCompletionsResponse(?callable $sseCallback = null): ChatCompletionsResponse
    {
        $isStream = !empty($sseCallback);
        $this->withStream($isStream);

        $cred = new Credential(
            $this->getKeyModel()->config['secret_id'],
            $this->getKeyModel()->config['secret_key']
        );
        $httpProfile = new HttpProfile();
        $httpProfile->setEndpoint($this->host);

        $clientProfile = new ClientProfile();
        $clientProfile->setHttpProfile($httpProfile);

        $client = new HunyuanClient($cred, "", $clientProfile);
        $isStream && $client->setSseResponseCallbackFunc($sseCallback);

        $req = new ChatCompletionsRequest();
        $req->fromJsonString(json_encode($this->getPayload()));

        $res = $client->ChatCompletions($req);
        if (!empty($res->getErrorMsg())) {
            throw new CompletionException(
                $res->getErrorMsg()->getMsg(),
                $res->getErrorMsg()->getCode()
            );
        }
        return $res;
    }
}